#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Milvus集合删除脚本
用于删除所有或指定的Milvus collections
"""
import sys
import os
from pathlib import Path
import argparse

# 添加backend目录到系统路径（仅用于测试脚本）
backend_dir = Path(__file__).parent.parent.parent
sys.path.insert(0, str(backend_dir))

try:
    # 尝试安全导入
    from pymilvus import connections, utility
except ImportError as e:
    print(f"错误: 导入pymilvus失败 - {e}")
    print("请执行: pip install -r backend/requirements.txt")
    sys.exit(1)

# 导入配置管理器
try:
    from app.utils.config_manager import config
except ImportError as e:
    print(f"错误: 导入配置管理器失败 - {e}")
    print("请确保config_manager.py文件存在于正确的路径")
    sys.exit(1)


def connect_to_milvus():
    """连接到Milvus服务器"""
    milvus_uri = config.milvus_config.get("uri")
    milvus_token = config.milvus_config.get("token")
    
    if not milvus_uri:
        print("错误: 未配置Milvus URI，请检查配置文件")
        return False
    
    try:
        if milvus_token:
            connections.connect("default", uri=milvus_uri, token=milvus_token)
        else:
            connections.connect("default", uri=milvus_uri)
        print(f"成功连接到Milvus服务器: {milvus_uri}")
        return True
    except Exception as e:
        print(f"连接Milvus服务器失败: {e}")
        return False


def get_all_collections():
    """获取所有集合名称"""
    try:
        collections = utility.list_collections()
        return collections
    except Exception as e:
        print(f"获取集合列表失败: {e}")
        return []


def drop_collection(collection_name):
    """删除指定集合"""
    try:
        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)
            print(f"成功删除集合: {collection_name}")
            return True
        else:
            print(f"集合不存在: {collection_name}")
            return False
    except Exception as e:
        print(f"删除集合 {collection_name} 失败: {e}")
        return False


def confirm_action(message):
    """确认用户操作"""
    response = input(f"{message} (y/n): ").lower().strip()
    return response == 'y' or response == 'yes'


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Milvus集合删除工具")
    parser.add_argument("-a", "--all", action="store_true", help="删除所有集合")
    parser.add_argument("-c", "--collection", help="指定要删除的集合名称")
    parser.add_argument("-f", "--force", action="store_true", help="强制删除，不需要确认")
    parser.add_argument("-y", "--yes", action="store_true", help="对所有提示回答是")
    args = parser.parse_args()
    
    # 获取配置信息
    try:
        config_path = config.get_config_file_path()
        print(f"使用配置文件: {config_path}")
    except Exception as e:
        print(f"获取配置文件路径失败: {e}")
    
    # 连接Milvus
    if not connect_to_milvus():
        return
    
    # 获取所有集合
    all_collections = get_all_collections()
    
    if not all_collections:
        print("Milvus中没有找到任何集合")
        return
    
    print(f"Milvus中共有 {len(all_collections)} 个集合:")
    for i, name in enumerate(all_collections, 1):
        print(f"{i}. {name}")
    
    # 删除指定集合
    if args.collection:
        if args.collection not in all_collections:
            print(f"错误: 指定的集合 '{args.collection}' 不存在")
            return
        
        if args.force or args.yes or confirm_action(f"确定要删除集合 '{args.collection}' 吗?"):
            drop_collection(args.collection)
    
    # 删除所有集合
    elif args.all:
        if args.force or args.yes or confirm_action(f"警告: 确定要删除所有 {len(all_collections)} 个集合吗?"):
            for collection_name in all_collections:
                drop_collection(collection_name)
            print("所有集合删除操作已完成")
    
    # 交互模式
    else:
        print("\n请选择操作:")
        print("1. 删除所有集合")
        print("2. 删除特定集合")
        print("3. 退出")
        
        choice = input("请输入选项 (1-3): ").strip()
        
        if choice == "1":
            if args.yes or confirm_action(f"警告: 确定要删除所有 {len(all_collections)} 个集合吗?"):
                for collection_name in all_collections:
                    drop_collection(collection_name)
                print("所有集合删除操作已完成")
        
        elif choice == "2":
            collection_index = input(f"请输入要删除的集合编号 (1-{len(all_collections)}): ").strip()
            try:
                idx = int(collection_index) - 1
                if 0 <= idx < len(all_collections):
                    collection_name = all_collections[idx]
                    if args.yes or confirm_action(f"确定要删除集合 '{collection_name}' 吗?"):
                        drop_collection(collection_name)
                else:
                    print("无效的集合编号")
            except ValueError:
                print("请输入有效的数字")
        
        elif choice == "3":
            print("操作已取消")
        
        else:
            print("无效的选项")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n操作已取消")
    except Exception as e:
        print(f"发生未处理的错误: {e}")
        import traceback
        traceback.print_exc() 